#include <PfRingDeviceList.h>
#include <Packet.h>
#include <iostream>
#include <getopt.h>
#include "headers/packet_handler.h"
#include "headers/PacketStats.h"
#include <IPv4Layer.h>
#include <PcapFileDevice.h>
#include <chrono>
#include <signal.h>
#include <atomic>
#include <thread>
#include <clickhouse/client.h>
#include <csignal>
#include <pcap.h>


#define EXIT_WITH_ERROR(reason)                       \
	do                                                \
	{                                                 \
		std::cout << std::endl                        \
				  << "ERROR: " << reason << std::endl \
				  << std::endl;                       \
		exit(1);                                      \
	} while (0)

#define numOfCaptureThreads 12
using namespace clickhouse;

std::atomic<long long> packet_count;

void onPacketArrives(pcpp::RawPacket *packet, uint32_t numOfPackets, uint8_t threadId, pcpp::PfRingDevice *dev, void *cookie)
{
	packet_count++;

	pcpp::Packet parsedPacket(packet);
	static thread_local packet_handler local_handler; // thread-local instance
	if (parsedPacket.isPacketOfType(pcpp::IPv4))
	{
		local_handler.handle_packet(&parsedPacket);
		auto *stats = static_cast<PacketStats *>(cookie);

		stats->consumePacketProtocol(parsedPacket);
		stats->updateSessionStats(parsedPacket,packet_count);
		stats->consumePacket5minute(parsedPacket);
		stats->consumePacketPortTraffic(parsedPacket);
	}
}

pcpp::PfRingDevice *init(std::string interfaceNameOrIP = "enp7s0")
{
	// 获取网卡设备
	pcpp::PfRingDevice *device = pcpp::PfRingDeviceList::getInstance().getPfRingDeviceByName(interfaceNameOrIP);

	if (!device->openMultiRxChannels(numOfCaptureThreads, pcpp::PfRingDevice::RoundRobin))
		EXIT_WITH_ERROR("Couldn't open " << numOfCaptureThreads << " RX channels on interface '" << device->getDeviceName() << "'");

	int totalNumOfCores = pcpp::getNumOfCores();
	std::cout << "当前设备总核心数:" << totalNumOfCores << std::endl;
	std::cout << "设置的捕获线程数:" << numOfCaptureThreads << std::endl;
	try {
		/// 如果不存在表,就创建表
		ClientOptions options = ClientOptions().SetHost("localhost");
		Client client(options);
		client.Execute(R"(CREATE TABLE IF NOT EXISTS default._packets (insertTime DateTime DEFAULT now(),hashVal UInt32,packetLen Int32,tvSec UInt32,tvNsec UInt32,srcIp String,
					dstIp String,ttl UInt8,tos UInt8,id UInt16,offset UInt16,protocolIP  UInt8,protocol UInt8,lenLoad UInt16,srcPort UInt16,
					dstPort UInt16,ackNum UInt32,seqNum UInt32,flag UInt16,window UInt16,lenUdp UInt16,icmpType UInt8,icmpCode UInt8,) 
				ENGINE = MergeTree()
				ORDER BY (hashVal,insertTime)
				TTL insertTime + INTERVAL 2 HOUR DELETE
		)");
		client.Execute(R"(
			CREATE TABLE IF NOT EXISTS default._protocol_hash_table (
				protocol UInt8,
				hashVal UInt32
			)
			ENGINE = MergeTree()
			ORDER BY (protocol)
		)");
		client.Execute(R"(
			CREATE TABLE IF NOT EXISTS default._src_ip_hash_table (
				srcIp String,
				hashVal UInt32
			)
			ENGINE = MergeTree()
			ORDER BY (srcIp)
		)");

		// 创建一个表,用来对结果进行聚合，存储default.packets表中每个hashVal的最近更新时间，并统计每个hashVal的行数(packet_count),
		client.Execute(R"(
			CREATE TABLE IF NOT EXISTS default._hashVal_summary(
				hashVal UInt32,
				latest_insertTime AggregateFunction(max, DateTime),
				packet_count AggregateFunction(count, Nullable(Int32)),
				is_infered Bool
			)
			ENGINE = AggregatingMergeTree()
			ORDER BY hashVal;
		)");

		// 创建一个物化视图hashVal_mv，存储default.packets表中每个hashVal的最近更新时间，并统计每个hashVal的行数(packet_count),
		client.Execute(R"(
			CREATE MATERIALIZED VIEW IF NOT EXISTS default._hashVal_mv TO default._hashVal_summary
			AS SELECT
				hashVal,
				maxState(insertTime) as latest_insertTime,
				countState() as packet_count,
				false as is_infered
			FROM default._packets
			GROUP BY hashVal;
		)");
	} catch (const std::exception& e) {
		// 处理标准异常
		std::cerr << "Standard exception: " << e.what() << std::endl;
		throw std::runtime_error("Standard exception caught in live_uni L115");
	} catch (...) {
		// 处理其他未知异常
		std::cerr << "Unknown exception occurred" << std::endl;
		throw(1);
	}

	return device;
}

void startCapture(pcpp::PfRingDevice *device, PacketStats *stats)
{
	pcpp::CoreMask coreMask = 0;
	int threadId = 0;
	int threadCount = 0;

	// coreMask update
	while (threadCount < numOfCaptureThreads)
	{
		// 当前的线程ID是否与设备当前核心ID相同
		if (pcpp::SystemCores::IdToSystemCore[threadId].Id != device->getCurrentCoreId().Id)
		{
			coreMask |= pcpp::SystemCores::IdToSystemCore[threadId].Mask; // 如果不同，则使用该核心，从而提高使用效率
			threadCount++;
		}

		threadId++;
	}

	std::cout << "开始捕获流量.." << std::endl;

	if (!device->startCaptureMultiThread(onPacketArrives, stats, coreMask))
	{
		EXIT_WITH_ERROR("Couldn't start capturing on core mask 0x" << std::hex << coreMask << " on interface '" << device->getDeviceName() << "'");
	}
	return;
}
void stopCapture(pcpp::PfRingDevice *device)
{
	// 停止捕获
	device->stopCapture();

	// 关闭设备
	device->close();
	std::cout << "停止捕获流量..." << std::endl;
}

// 全局变量注册
pcpp::PfRingDevice *dev;
volatile sig_atomic_t is_stop = 0;

void signal_handler(int signum)
{
	stopCapture(dev);
	std::cout << "count数量" << packet_count << std::endl;
	is_stop = 1;
}

int main(int argc, char *argv[])
{
	signal(SIGINT, signal_handler);

	// 初始化
	dev = init("enp7s0");
	PacketStats stats;
	
	startCapture(dev, &stats);
	std::thread statsPrinter([&]()
	{
        while (!is_stop)
        {
            std::this_thread::sleep_for(std::chrono::seconds(12)); // 每12秒打印一次统计信息
			//调用函数历史记录文件写入
			stats.out_5min_file();
			stats.out_protocol_file();
			stats.out_port_file();
        } 
	});

	while (!is_stop)
	{
		std::this_thread::sleep_for(std::chrono::seconds(30));
	};

	statsPrinter.join(); // 等待统计打印线程结束

	// Clear stats
	stats.clear();

	return 0;
}